using UnityEngine;

public class Fractal : MonoBehaviour
{
    private struct FractalPart
    {
        public Vector3 direction, worldPosition;
        public Quaternion rotation, worldRotation;
        public float spinAngle;
    }
    
    [SerializeField, Range(1, 8)]
    private int depth = 4;
    [SerializeField]
    private Mesh mesh;
    [SerializeField]
    private Material material;

    private static Vector3[] _directions = { Vector3.up, Vector3.right, Vector3.left, Vector3.forward, Vector3.back };
    private static Quaternion[] _rotations =
    {
        Quaternion.identity, Quaternion.Euler(0f, 0f, -90f), Quaternion.Euler(0f, 0f, 90f),
        Quaternion.Euler(90f, 0f, 0f), Quaternion.Euler(-90f, 0f, 0f)
    };

    private static int _matricesId = Shader.PropertyToID("_Matrices");
    private static MaterialPropertyBlock propertyBlock;

    private FractalPart[][] _parts;
    private Matrix4x4[][] _matrices;
    private ComputeBuffer[] _matricesBuffers;

    private void OnEnable()
    {
        _parts = new FractalPart[depth][];
        _matrices = new Matrix4x4[depth][];
        _matricesBuffers = new ComputeBuffer[depth];
        int stride = 16 * 4;
        for (int i = 0, size = 1; i < _parts.Length; i++, size *= 5)
        {
            _parts[i] = new FractalPart[size];
            _matrices[i] = new Matrix4x4[size];
            _matricesBuffers[i] = new ComputeBuffer(size, stride);
        }

        _parts[0][0] = CreatePart(0);
        for (int i = 1; i < _parts.Length; i++)
        {
            var levelParts = _parts[i];
            for (int j = 0; j < levelParts.Length; j += 5)
            {
                for (int k = 0; k < 5; k++)
                {
                    levelParts[j + k] = CreatePart(k);
                }
            }
        }

        if (propertyBlock == null)
            propertyBlock = new MaterialPropertyBlock();
    }

    private void OnDisable()
    {
        for (int i = 0; i < _matricesBuffers.Length; i++)
        {
            _matricesBuffers[i].Release();
        }

        _parts = null;
        _matrices = null;
        _matricesBuffers = null;
    }

    private void OnValidate()
    {
        if (_parts != null && enabled)
        {
            OnDisable();
            OnEnable();
        }
    }

    private void Update()
    {
        var spinAngleDelta = 22.5f * Time.deltaTime;
        var rootPart = _parts[0][0];
        rootPart.spinAngle += spinAngleDelta;
        var rootTransform = transform;
        rootPart.worldRotation =
            rootTransform.rotation * (rootPart.rotation * Quaternion.Euler(0f, rootPart.spinAngle, 0f));
        rootPart.worldPosition = transform.position;
        _parts[0][0] = rootPart;

        float objectScale = rootTransform.lossyScale.x;
        _matrices[0][0] = Matrix4x4.TRS(rootPart.worldPosition, rootPart.worldRotation, objectScale * Vector3.one);

        var scale = objectScale;
        for (int i = 1; i < _parts.Length; i++)
        {
            scale *= 0.5f;
            var parentParts = _parts[i - 1];
            var levelParts = _parts[i];
            var levelMatrices = _matrices[i];
            for (int j = 0; j < levelParts.Length; j++)
            {
                var parent = parentParts[j / 5];
                var part = levelParts[j];
                part.spinAngle += spinAngleDelta;
                part.worldRotation = parent.worldRotation * (part.rotation * Quaternion.Euler(0f, part.spinAngle, 0f));
                part.worldPosition = parent.worldPosition + parent.worldRotation * (1.5f * scale * part.direction);
                levelParts[j] = part;
                levelMatrices[j] = Matrix4x4.TRS(part.worldPosition, part.worldRotation, scale * Vector3.one);
            }
        }

        var bounds = new Bounds(rootPart.worldPosition, 3f * objectScale * Vector3.one);
        for (int i = 0; i < _matricesBuffers.Length; i++)
        {
            var buffer = _matricesBuffers[i];
            buffer.SetData(_matrices[i]);
            propertyBlock.SetBuffer(_matricesId, buffer);
            Graphics.DrawMeshInstancedProcedural(mesh, 0, material, bounds, buffer.count, propertyBlock);
        }
    }

    private FractalPart CreatePart(int childIndex)
    {
        return new FractalPart
            {direction = _directions[childIndex], rotation = _rotations[childIndex]};
    }
}
